#!/usr/bin/env python

from Bio import SeqIO
from pathlib import Path
import sys
import argparse
import os

parser = argparse.ArgumentParser(description='This script will split a multifasta into individual fasta files, each file named with the header of the sequence within it, written to a new subdirectory. (It expects standard characters in the headers only, e.g. no spaces or special characters). For version info, run `bit-version`.')

required = parser.add_argument_group('required arguments')

required.add_argument("-i", "--input_fasta", help="Starting fasta file", action="store", dest="input_fasta", required=True)
parser.add_argument("-d", "--desired_subdirectory_name", help='Name of new subdirectory holding split sequences (default: "sub")', action="store", dest="subdirectory", default="sub")

if len(sys.argv)==1:
    parser.print_help(sys.stderr)
    sys.exit(0)

args = parser.parse_args()

out_dir = Path(str(args.subdirectory))

Path.mkdir(out_dir, parents=True, exist_ok=True)

in_fasta = open(args.input_fasta, "r")

for seq_record in SeqIO.parse(in_fasta, "fasta"):
    curr_header = str(seq_record.id)

    curr_out = open(str(args.subdirectory) + "/" + curr_header + ".fa", "w")
    curr_out.write(">" + str(seq_record.id) + "\n" + str(seq_record.seq) + "\n")
    curr_out.close()

in_fasta.close()
